{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "a8b892c8",
   "metadata": {},
   "source": [
    "Printed and electronic copies of *Modeling and Simulation in Python* are available from [No Starch Press](https://nostarch.com/modeling-and-simulation-python) and [Bookshop.org](https://bookshop.org/p/books/modeling-and-simulation-in-python-allen-b-downey/17836697?ean=9781718502161) and [Amazon](https://amzn.to/3y9UxNb)."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "colonial-jacob",
   "metadata": {},
   "source": [
    "# Sweeping Parameters"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "modified-discretion",
   "metadata": {
    "tags": []
   },
   "source": [
    "*Modeling and Simulation in Python*\n",
    "\n",
    "Copyright 2021 Allen Downey\n",
    "\n",
    "License: [Creative Commons Attribution-NonCommercial-ShareAlike 4.0 International](https://creativecommons.org/licenses/by-nc-sa/4.0/)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "complicated-retreat",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# install Pint if necessary\n",
    "\n",
    "try:\n",
    "    import pint\n",
    "except ImportError:\n",
    "    !pip install pint"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "selected-there",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# download modsim.py if necessary\n",
    "\n",
    "from os.path import basename, exists\n",
    "\n",
    "def download(url):\n",
    "    filename = basename(url)\n",
    "    if not exists(filename):\n",
    "        from urllib.request import urlretrieve\n",
    "        local, _ = urlretrieve(url, filename)\n",
    "        print('Downloaded ' + local)\n",
    "    \n",
    "download('https://raw.githubusercontent.com/AllenDowney/' +\n",
    "         'ModSimPy/master/modsim.py')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "prepared-apparatus",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# import functions from modsim\n",
    "\n",
    "from modsim import *"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "awful-martial",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "download('https://github.com/AllenDowney/ModSimPy/raw/master/chap11.py')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "changing-burner",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "download('https://github.com/AllenDowney/ModSimPy/raw/master/chap12.py')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "defensive-automation",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# import code from previous notebooks\n",
    "\n",
    "from chap11 import make_system\n",
    "from chap11 import update_func\n",
    "from chap11 import run_simulation\n",
    "from chap11 import plot_results\n",
    "\n",
    "from chap12 import calc_total_infected"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "structural-application",
   "metadata": {},
   "source": [
    "In the previous chapter we extended the Kermack-McKendrick (KM) model to include immunization and used it to demonstrate herd immunity.\n",
    "\n",
    "But we assumed that the parameters of the model, contact rate and\n",
    "recovery rate, were known. In this chapter, we'll explore the behavior of\n",
    "the model as we vary these parameters.\n",
    "\n",
    "In the next chapter, we'll use analysis to understand these relationships better, and propose a method for using data to estimate parameters."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "tracked-outreach",
   "metadata": {},
   "source": [
    "## Sweeping Beta\n",
    "\n",
    "Recall that $\\beta$ is the contact rate, which captures both the\n",
    "frequency of interaction between people and the fraction of those\n",
    "interactions that result in a new infection. If $N$ is the size of the\n",
    "population and $s$ is the fraction that's susceptible, $s N$ is the\n",
    "number of susceptibles, $\\beta s N$ is the number of contacts per day\n",
    "between susceptibles and other people, and $\\beta s i N$ is the number\n",
    "of those contacts where the other person is infectious.\n",
    "\n",
    "As $\\beta$ increases, we expect the total number of infections to\n",
    "increase. To quantify that relationship, I'll create a range of values\n",
    "for $\\beta$:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "forward-responsibility",
   "metadata": {},
   "outputs": [],
   "source": [
    "beta_array = linspace(0.1, 1.1, 11)\n",
    "gamma = 0.25"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "eligible-shareware",
   "metadata": {},
   "source": [
    "We'll start with a single value for `gamma`, which is the recovery rate, that is, the fraction of infected people who recover per day.\n",
    "\n",
    "The following function takes `beta_array` and `gamma` as parameters.\n",
    "It runs the simulation for each value of `beta` and computes the same metric we used in the previous chapter, the fraction of the population that gets infected.\n",
    "\n",
    "The result is a `SweepSeries` that contains the values of `beta` and the corresponding metrics."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "dutch-chemistry",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "def sweep_beta(beta_array, gamma):\n",
    "    sweep = SweepSeries()\n",
    "    for beta in beta_array:\n",
    "        system = make_system(beta, gamma)\n",
    "        results = run_simulation(system, update_func)\n",
    "        sweep[beta] = calc_total_infected(results, system)\n",
    "    return sweep"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "mathematical-process",
   "metadata": {},
   "source": [
    "We can run `sweep_beta` like this:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "christian-singer",
   "metadata": {},
   "outputs": [],
   "source": [
    "infected_sweep = sweep_beta(beta_array, gamma)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "collectible-positive",
   "metadata": {},
   "source": [
    "Before we plot the results, I will use a formatted string literal, also called an *f-string* to assemble a label that includes the value of `gamma`:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "proprietary-credit",
   "metadata": {},
   "outputs": [],
   "source": [
    "label = f'gamma = {gamma}'\n",
    "label"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "invisible-bible",
   "metadata": {},
   "source": [
    "An f-string starts with the letter `f` followed by a string in single or double quotes. \n",
    "The string can contain any number of format specifiers in squiggly brackets, `{}`.\n",
    "When a variable name appears in a format specifier, it is replaced with the value of the variable.\n",
    "\n",
    "In this example, the value of `gamma` is `0.25`, so the value of `label` is `'gamma = 0.25'`.\n",
    "\n",
    "You can read more about f-strings at <https://docs.python.org/3/tutorial/inputoutput.html#tut-f-strings>.\n",
    "\n",
    "Now let's see the results."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "developmental-video",
   "metadata": {},
   "outputs": [],
   "source": [
    "infected_sweep.plot(label=label, color='C1')\n",
    "\n",
    "decorate(xlabel='Contact rate (beta)',\n",
    "         ylabel='Fraction infected')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "durable-isaac",
   "metadata": {},
   "source": [
    "Remember that this figure\n",
    "is a parameter sweep, not a time series, so the x-axis is the parameter\n",
    "`beta`, not time.\n",
    "\n",
    "When `beta` is small, the contact rate is low and the outbreak never\n",
    "really takes off; the total number of infected students is near zero. As\n",
    "`beta` increases, it reaches a threshold near 0.3 where the fraction of\n",
    "infected students increases quickly. When `beta` exceeds 0.5, more than\n",
    "80% of the population gets sick."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "described-jacksonville",
   "metadata": {},
   "source": [
    "## Sweeping Gamma\n",
    "\n",
    "Let's see what that looks like for a few different values of `gamma`.\n",
    "We'll use `linspace` to make an array of values:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "chief-greece",
   "metadata": {},
   "outputs": [],
   "source": [
    "gamma_array = linspace(0.1, 0.7, 4)\n",
    "gamma_array"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "medium-greek",
   "metadata": {},
   "source": [
    "And run `sweep_beta` for each value of `gamma`:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "thermal-pulse",
   "metadata": {},
   "outputs": [],
   "source": [
    "for gamma in gamma_array:\n",
    "    infected_sweep = sweep_beta(beta_array, gamma)\n",
    "    label = f'gamma = {gamma}'\n",
    "    infected_sweep.plot(label=label)\n",
    "    \n",
    "decorate(xlabel='Contact rate (beta)',\n",
    "         ylabel='Fraction infected')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "artistic-charge",
   "metadata": {},
   "source": [
    "When `gamma` is low, the recovery rate is low, which means people are infectious longer.\n",
    "In that case, even a low contact rate (`beta`) results in an epidemic.\n",
    "\n",
    "When `gamma` is high, `beta` has to be even higher to get things going."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "metropolitan-collect",
   "metadata": {},
   "source": [
    "## Using a SweepFrame\n",
    "\n",
    "In the previous section, we swept a range of values for `gamma`, and for\n",
    "each value of `gamma`, we swept a range of values for `beta`. This process is a\n",
    "*two-dimensional sweep*.\n",
    "\n",
    "If we want to store the results, rather than plot them, we can use a\n",
    "`SweepFrame`, which is a kind of `DataFrame` where the rows sweep one\n",
    "parameter, the columns sweep another parameter, and the values contain\n",
    "metrics from a simulation.\n",
    "\n",
    "This function shows how it works:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "scientific-dealer",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "def sweep_parameters(beta_array, gamma_array):\n",
    "    frame = SweepFrame(columns=gamma_array)\n",
    "    for gamma in gamma_array:\n",
    "        frame[gamma] = sweep_beta(beta_array, gamma)\n",
    "    return frame"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "educated-congress",
   "metadata": {},
   "source": [
    "`sweep_parameters` takes as parameters an array of values for `beta` and\n",
    "an array of values for `gamma`.\n",
    "\n",
    "It creates a `SweepFrame` to store the results, with one column for each\n",
    "value of `gamma` and one row for each value of `beta`.\n",
    "\n",
    "Each time through the loop, we run `sweep_beta`. The result is a\n",
    "`SweepSeries` object with one element for each value of `beta`. The\n",
    "assignment inside the loop stores the `SweepSeries` as a new column in\n",
    "the `SweepFrame`, corresponding to the current value of `gamma`.\n",
    "\n",
    "At the end, the `SweepFrame` stores the fraction of students infected\n",
    "for each pair of parameters, `beta` and `gamma`.\n",
    "\n",
    "We can run `sweep_parameters` like this:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "engaging-philosophy",
   "metadata": {},
   "outputs": [],
   "source": [
    "frame = sweep_parameters(beta_array, gamma_array)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "realistic-hughes",
   "metadata": {},
   "source": [
    "With the results in a `SweepFrame`, we can plot each column like this:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "id": "false-forestry",
   "metadata": {},
   "outputs": [],
   "source": [
    "for gamma in gamma_array:\n",
    "    label = f'gamma = {gamma}'\n",
    "    frame[gamma].plot(label=label)\n",
    "\n",
    "decorate(xlabel='Contact rate (beta)',\n",
    "         ylabel='Fraction infected',\n",
    "         title='Sweep beta, multiple values of gamma')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "armed-spending",
   "metadata": {},
   "source": [
    "Alternatively, we can plot each row like this:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "id": "silent-advantage",
   "metadata": {},
   "outputs": [],
   "source": [
    "for beta in [0.2, 0.5, 0.8, 1.1]:\n",
    "    label = f'beta = {beta}'\n",
    "    frame.loc[beta].plot(label=label)\n",
    "    \n",
    "decorate(xlabel='Recovery rate (gamma)',\n",
    "         ylabel='Fraction infected',\n",
    "         title='Sweep gamma, multiple values of beta')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "diverse-works",
   "metadata": {},
   "source": [
    "This example demonstrates one use of a `SweepFrame`: we can run the analysis once, save the results, and then generate different visualizations.\n",
    "\n",
    "Another way to visualize the results of a two-dimensional sweep is a\n",
    "*contour plot*, which shows the parameters on the axes and contour\n",
    "lines where the value of the metric is constant.\n",
    "\n",
    "The ModSim library provides `contour`, which takes a `SweepFrame` and uses Matplotlib to generate a contour plot."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "id": "veterinary-cheese",
   "metadata": {},
   "outputs": [],
   "source": [
    "contour(frame)\n",
    "\n",
    "decorate(xlabel='Recovery rate (gamma)',\n",
    "         ylabel='Contact rate (beta)',\n",
    "         title='Contour plot, fraction infected')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "operational-daughter",
   "metadata": {},
   "source": [
    "The values of `gamma` are on the $x$-axis, corresponding to the columns of the `SweepFrame`.\n",
    "The values of `beta` are on the $y$-axis, corresponding to the rows of the `SweepFrame`.\n",
    "Each line follows a contour where the infection rate is constant.\n",
    "\n",
    "Infection rates are lowest in the lower right, where the contact rate is low and the recovery rate is high. They increase as we move to the upper left, where the contact rate is high and the recovery rate is low.\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "personal-formula",
   "metadata": {},
   "source": [
    "## Summary\n",
    "\n",
    "This chapter demonstrates a two-dimensional parameter sweep using a `SweepFrame` to store the results.\n",
    "\n",
    "We plotted the results three ways: \n",
    "\n",
    "* First we plotted total infections versus `beta`, with one line for each value of `gamma`.\n",
    "\n",
    "* Then we plotted total infections versus `gamma`, with one line for each value of `beta`.\n",
    "\n",
    "* Finally, we made a contour plot with `beta` on the $y$-axis, `gamma` on the $x$-axis and contour lines where the metric is constant.\n",
    "\n",
    "These visualizations suggest that there is a relationship between `beta` and `gamma` that determines the outcome of the model. \n",
    "In fact, there is.\n",
    "In the next chapter we'll explore it by running simulations, then derive it by analysis."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "variable-capture",
   "metadata": {},
   "source": [
    "## Exercises\n",
    "\n",
    "This chapter is available as a Jupyter notebook where you can read the text, run the code, and work on the exercises. \n",
    "You can access the notebooks at <https://allendowney.github.io/ModSimPy/>."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "conscious-blues",
   "metadata": {},
   "source": [
    "### Exercise 1\n",
    "\n",
    "  If we know `beta` and `gamma`, we can compute the fraction of the population that gets infected.  In general, we don't know these parameters, but sometimes we can estimate them based on the behavior of an outbreak.\n",
    "\n",
    "Suppose the infectious period for the Freshman Plague is known to be 2 days on average, and suppose during one particularly bad year, 40% of the class is infected at some point.  Estimate the time between contacts, `1/beta`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "id": "democratic-driving",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Solution goes here"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "id": "stretch-israel",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Solution goes here"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "id": "controlling-strand",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Solution goes here"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "instructional-hayes",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "celltoolbar": "Tags",
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
